#include <openssl/bn.h>
#include "../../core/pch.h"
#include "../../pi/process.h"
#include "../../pi/binarysemaphore.h"

/* Structure for receiver/worker address. */
static struct sockaddr_in receiver_addr;
static struct sockaddr_in worker_addr;

/* Receiver socket. */
static int receiver_sock;

/* Structure for client address. */
static struct sockaddr_in cliaddr;

/* Algorithm specific. */
static ULL N, p, times_count, num;

/* Check if number is prime.
   Uses OpenSSL library. */

int 
isPrime(ULL x)
{
	char buf[128];
	memset(buf, 0, 128);
	sprintf(buf, "%llu", x);

	BIGNUM* bn_n = BN_new();
	BN_dec2bn(&bn_n, buf);
	int ret = BN_is_prime(bn_n, 0, NULL, NULL, NULL);
	BN_free(bn_n);

	return ret;
}

/* This function initializes network for
   receiver process. It creates socket for
   this process and starts listening on it
   on given by master process port. */

static void 
init_network(struct sockaddr_in &recv_addr, int &recv_sock, int recv_port)
{
	int ret;

	FILL_SOCKADDR_IN_INT(recv_addr, AF_INET, recv_port, INADDR_ANY);

	recv_sock = socket(AF_INET, SOCK_STREAM, 0); 
	if (recv_sock == SOCKET_ERROR)
	{
		DEBUGS("receiver: socket error (%s)", strerror(errno));
		exit(SOCKET_ERROR);
	}

	ret = bind(recv_sock, (struct sockaddr*) &recv_addr, sizeof(recv_addr)); 
	if (ret == BIND_ERROR)
	{
		DEBUGS("receiver: bind error (%s)", strerror(errno));
		exit(BIND_ERROR);
	}

	ret = listen(recv_sock, 64); 
	if (ret == LISTEN_ERROR)
	{
		DEBUGS("receiver: listen error (%s)", strerror(errno));
		exit(LISTEN_ERROR);
	}
}

/* Send to all slaves END_OF_COMPUTATION packet.
   After receiving this packet slaves should kill
   themselves. */

static void
kill_slaves(struct info_table *shm)
{
	for (int i = 0; i < shm->count; i++)
	{
		FILL_SOCKADDR_IN_TXT(worker_addr, AF_INET, 
			shm->worker_port[i], shm->ip[i]);

		int worker_sock = socket(AF_INET, SOCK_STREAM, 0);
		if (worker_sock == SOCKET_ERROR)
		{
			DEBUGS("sender: socket error for worker (%s)", strerror(errno));
			exit(SOCKET_ERROR);
		}

		int ret = connect(worker_sock, (struct sockaddr*) &worker_addr, sizeof(worker_addr));
		if (ret == CONNECT_ERROR)
		{
			DEBUGS("sender: connect error to worker (%s)", strerror(errno));
			close(worker_sock);
			exit(CONNECT_ERROR);
		}

		char msg = END_OF_COMPUTATION;
		send(worker_sock, &msg, sizeof(char), 0);
		close(worker_sock);
	}
}

/* Write result to file. */

void
write_result(FILE *out, ULL n, ULL p, UINT times_count)
{
	fprintf(out, "N = %llu, p = %llu, times_count = %d\n", n, p, times_count); 
	fflush(out);
}

/* This function cleans after factorization is
   ended. Kills other processes, writes final results. */

static void
finalize(FILE *out, struct info_table *shm, int &sender_pid,
	int &logger_pid)
{
	/* Kill all slaves. */
	kill_slaves(shm);

	/* Write final result. */
	write_result(out, shm->n, shm->n, 1);

	/* Kill other master processes. */
	kill(sender_pid, SIGKILL);
	kill(logger_pid, SIGKILL);
}

/* Main receiver function starts listening for results.
   When it accepts a result it divides current number
   in seperate process. */

int receiver(int segment_id, int sem_id, int receiver_port, struct info_table *shm,
	FILE *out_file, FILE *primes_file, int &sender_pid, int &logger_pid)
{
	DEBUGS("receiver: running");

	/* Initialize network stuff. */
	init_network(receiver_addr, receiver_sock, receiver_port);

	/* Start listening for incoming connections. */
	while (true)
	{
		DEBUGS("receiver: waiting for results from slaves");

		/* Accept connection from slave. */
		socklen_t len = sizeof(struct sockaddr);
		int worker_sock = accept(receiver_sock, (struct sockaddr*) &cliaddr, &len);
		if (worker_sock == ACCEPT_ERROR)
		{
			DEBUGS("receiver: accept error (%s)", strerror(errno));
			exit(ACCEPT_ERROR);
		}

		/* Receive results from slave. */
		recv(worker_sock, &N, sizeof(unsigned long long), 0);
		recv(worker_sock, &p, sizeof(unsigned long long), 0);
		recv(worker_sock, &times_count, sizeof(unsigned int), 0);
		recv(worker_sock, &num, sizeof(unsigned long long), 0);

		DEBUGS("receiver: accepted connection from slave, results for %llu and prime %llu", N, p);

		/* Run process to use computation results. */
		if (shm->n % p == 0)
		{
			if (Process::Create() == 0)
			{
				BinarySemaphore::Decrement(sem_id);

				/* Write results to file. */
				if (times_count > 0)
					write_result(out_file, shm->n, p, times_count); 

				/* Divide our current number by received factor. */
				while (times_count-- > 0) 
					shm->n = shm->n / p;		

				/* If after division our number is prime or it
				   is equal to one, we send kill signals to all
				   slaves and sender, logger. */
				if (isPrime(shm->n) || shm->n == 1)
					finalize(out_file, shm, sender_pid, logger_pid);

				BinarySemaphore::Increment(sem_id);

				/* Close files. */
				fclose(out_file);
				fclose(primes_file);

				exit(STATUS_OK);
			}
		}

		close(worker_sock);
	}

	return STATUS_OK;
}


